"""
优化的JSON序列化器实现
使用高性能JSON库和压缩技术
"""

import time
import uuid
import zlib
import gzip
from typing import Dict, Any, Optional, Union, Callable
from dataclasses import dataclass, asdict
from datetime import datetime, date
from decimal import Decimal
from io import StringIO, BytesIO

# 尝试导入高性能JSON库
JSON_LIBS = {}

try:
    import orjson

    JSON_LIBS['orjson'] = orjson
except ImportError:
    pass

try:
    import ujson

    JSON_LIBS['ujson'] = ujson
except ImportError:
    pass

# 标准库JSON作为后备
import json

JSON_LIBS['json'] = json

# 选择最优的JSON库
if 'orjson' in JSON_LIBS:
    DEFAULT_JSON_LIB = 'orjson'
elif 'ujson' in JSON_LIBS:
    DEFAULT_JSON_LIB = 'ujson'
else:
    DEFAULT_JSON_LIB = 'json'

print(f"使用JSON库: {DEFAULT_JSON_LIB}")


@dataclass
class JsonStats:
    """JSON序列化统计信息"""
    total_serialized: int = 0
    total_deserialized: int = 0
    total_bytes: int = 0
    compressed_bytes: int = 0
    avg_serialize_time: float = 0.0
    avg_deserialize_time: float = 0.0
    compression_ratio: float = 0.0
    error_count: int = 0


class JsonEncoder:
    """自定义JSON编码器"""

    def __init__(self):
        self.custom_encoders = {}
        self._register_default_encoders()

    def _register_default_encoders(self):
        """注册默认的编码器"""
        self.custom_encoders[datetime] = lambda dt: {
            '__type__': 'datetime',
            '__value__': dt.isoformat()
        }

        self.custom_encoders[date] = lambda d: {
            '__type__': 'date',
            '__value__': d.isoformat()
        }

        self.custom_encoders[Decimal] = lambda d: {
            '__type__': 'decimal',
            '__value__': str(d)
        }

        self.custom_encoders[bytes] = lambda b: {
            '__type__': 'bytes',
            '__value__': b.hex()
        }

        self.custom_encoders[set] = lambda s: {
            '__type__': 'set',
            '__value__': list(s)
        }

        self.custom_encoders[frozenset] = lambda fs: {
            '__type__': 'frozenset',
            '__value__': list(fs)
        }

    def register_encoder(self, type_class, encoder_func):
        """注册自定义编码器"""
        self.custom_encoders[type_class] = encoder_func

    def encode_object(self, obj):
        """编码对象"""
        obj_type = type(obj)

        if obj_type in self.custom_encoders:
            return self.custom_encoders[obj_type](obj)

        # 处理自定义对象
        if hasattr(obj, '__dict__'):
            return {
                '__type__': 'object',
                '__class__': obj.__class__.__name__,
                '__module__': obj.__class__.__module__,
                '__data__': obj.__dict__
            }

        raise TypeError(f"无法序列化类型: {obj_type}")


class JsonDecoder:
    """自定义JSON解码器"""

    def __init__(self):
        self.custom_decoders = {}
        self._register_default_decoders()

    def _register_default_decoders(self):
        """注册默认的解码器"""
        self.custom_decoders['datetime'] = lambda v: datetime.fromisoformat(v)
        self.custom_decoders['date'] = lambda v: date.fromisoformat(v)
        self.custom_decoders['decimal'] = lambda v: Decimal(v)
        self.custom_decoders['bytes'] = lambda v: bytes.fromhex(v)
        self.custom_decoders['set'] = lambda v: set(v)
        self.custom_decoders['frozenset'] = lambda v: frozenset(v)

    def register_decoder(self, type_name, decoder_func):
        """注册自定义解码器"""
        self.custom_decoders[type_name] = decoder_func

    def decode_object(self, obj):
        """解码对象"""
        if isinstance(obj, dict) and '__type__' in obj:
            type_name = obj['__type__']

            if type_name in self.custom_decoders:
                return self.custom_decoders[type_name](obj['__value__'])
            elif type_name == 'object':
                # 重建自定义对象
                try:
                    module = __import__(obj['__module__'], fromlist=[obj['__class__']])
                    cls = getattr(module, obj['__class__'])
                    instance = cls.__new__(cls)
                    instance.__dict__.update(self._decode_recursive(obj['__data__']))
                    return instance
                except (ImportError, AttributeError):
                    # 如果无法重建，返回原始字典
                    return obj

        return self._decode_recursive(obj)

    def _decode_recursive(self, obj):
        """递归解码对象"""
        if isinstance(obj, dict):
            return {k: self.decode_object(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [self.decode_object(item) for item in obj]
        else:
            return obj


class OptimizedJsonSerializer:
    """优化的JSON序列化器"""

    def __init__(self, json_lib: str = None, use_compression: bool = True,
                 compression_algorithm: str = 'gzip', compression_threshold: int = 1024,
                 ensure_ascii: bool = False, indent: Optional[int] = None):

        # 选择JSON库
        self.json_lib = json_lib or DEFAULT_JSON_LIB
        if self.json_lib not in JSON_LIBS:
            raise ValueError(f"不支持的JSON库: {self.json_lib}")

        self.json_module = JSON_LIBS[self.json_lib]
        self.use_compression = use_compression
        self.compression_algorithm = compression_algorithm
        self.compression_threshold = compression_threshold
        self.ensure_ascii = ensure_ascii
        self.indent = indent

        # 统计信息
        self.stats = JsonStats()

        # 编码器和解码器
        self.encoder = JsonEncoder()
        self.decoder = JsonDecoder()

    def _compress_data(self, data: bytes) -> bytes:
        """压缩数据"""
        if not self.use_compression or len(data) < self.compression_threshold:
            return data

        if self.compression_algorithm == 'gzip':
            return gzip.compress(data, compresslevel=6)
        elif self.compression_algorithm == 'zlib':
            return zlib.compress(data, level=6)
        else:
            return data

    def _decompress_data(self, data: bytes) -> bytes:
        """解压缩数据"""
        if not self.use_compression:
            return data

        try:
            if self.compression_algorithm == 'gzip':
                return gzip.decompress(data)
            elif self.compression_algorithm == 'zlib':
                return zlib.decompress(data)
        except Exception:
            # 如果解压缩失败，可能是未压缩的数据
            return data

        return data

    def dumps(self, obj: Any) -> bytes:
        """序列化对象为JSON字节流"""
        start_time = time.perf_counter()

        try:
            # 预处理对象
            processed_obj = self._preprocess_object(obj)

            # JSON序列化
            if self.json_lib == 'orjson':
                # orjson直接返回bytes
                json_data = self.json_module.dumps(
                    processed_obj,
                    option=orjson.OPT_UTC_Z if not self.ensure_ascii else 0
                )
            elif self.json_lib == 'ujson':
                # ujson返回字符串
                json_str = self.json_module.dumps(
                    processed_obj,
                    ensure_ascii=self.ensure_ascii,
                    indent=self.indent
                )
                json_data = json_str.encode('utf-8')
            else:
                # 标准库json
                json_str = self.json_module.dumps(
                    processed_obj,
                    ensure_ascii=self.ensure_ascii,
                    indent=self.indent,
                    default=self.encoder.encode_object
                )
                json_data = json_str.encode('utf-8')

            # 可选压缩
            original_size = len(json_data)
            compressed_data = self._compress_data(json_data)

            # 添加压缩标记
            if len(compressed_data) < original_size:
                result = b'\x01' + compressed_data  # 0x01表示已压缩
                self.stats.compressed_bytes += len(result)
            else:
                result = b'\x00' + json_data  # 0x00表示未压缩

            # 更新统计
            self.stats.total_serialized += 1
            self.stats.total_bytes += len(result)

            return result

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_serialize_time(duration)

    def loads(self, data: bytes) -> Any:
        """从JSON字节流反序列化对象"""
        start_time = time.perf_counter()

        try:
            if len(data) < 1:
                raise ValueError("数据太短")

            # 检查压缩标记
            compression_flag = data[0]
            payload = data[1:]

            # 解压缩（如果需要）
            if compression_flag == 0x01:
                payload = self._decompress_data(payload)

            # JSON反序列化
            if self.json_lib == 'orjson':
                obj = self.json_module.loads(payload)
            elif self.json_lib == 'ujson':
                json_str = payload.decode('utf-8')
                obj = self.json_module.loads(json_str)
            else:
                json_str = payload.decode('utf-8')
                obj = self.json_module.loads(json_str)

            # 后处理对象
            result = self.decoder.decode_object(obj)

            # 更新统计
            self.stats.total_deserialized += 1

            return result

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_deserialize_time(duration)

    def stream_dumps(self, objects: list) -> bytes:
        """流式序列化多个对象"""
        start_time = time.perf_counter()

        try:
            # 使用JSON Lines格式
            json_lines = []

            for obj in objects:
                processed_obj = self._preprocess_object(obj)

                if self.json_lib == 'orjson':
                    json_line = self.json_module.dumps(processed_obj)
                elif self.json_lib == 'ujson':
                    json_str = self.json_module.dumps(processed_obj, ensure_ascii=self.ensure_ascii)
                    json_line = json_str.encode('utf-8')
                else:
                    json_str = self.json_module.dumps(
                        processed_obj,
                        ensure_ascii=self.ensure_ascii,
                        default=self.encoder.encode_object
                    )
                    json_line = json_str.encode('utf-8')

                json_lines.append(json_line)

            # 合并所有行
            json_data = b'\n'.join(json_lines)

            # 压缩
            compressed_data = self._compress_data(json_data)

            if len(compressed_data) < len(json_data):
                result = b'\x01' + compressed_data
                self.stats.compressed_bytes += len(result)
            else:
                result = b'\x00' + json_data

            # 更新统计
            self.stats.total_serialized += len(objects)
            self.stats.total_bytes += len(result)

            return result

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_serialize_time(duration)

    def stream_loads(self, data: bytes) -> list:
        """流式反序列化多个对象"""
        start_time = time.perf_counter()

        try:
            if len(data) < 1:
                raise ValueError("数据太短")

            # 检查压缩标记
            compression_flag = data[0]
            payload = data[1:]

            # 解压缩（如果需要）
            if compression_flag == 0x01:
                payload = self._decompress_data(payload)

            # 按行分割
            json_lines = payload.split(b'\n')
            objects = []

            for json_line in json_lines:
                if not json_line.strip():
                    continue

                if self.json_lib == 'orjson':
                    obj = self.json_module.loads(json_line)
                elif self.json_lib == 'ujson':
                    json_str = json_line.decode('utf-8')
                    obj = self.json_module.loads(json_str)
                else:
                    json_str = json_line.decode('utf-8')
                    obj = self.json_module.loads(json_str)

                result = self.decoder.decode_object(obj)
                objects.append(result)

            # 更新统计
            self.stats.total_deserialized += len(objects)

            return objects

        except Exception as e:
            self.stats.error_count += 1
            raise
        finally:
            duration = time.perf_counter() - start_time
            self._update_deserialize_time(duration)

    def _preprocess_object(self, obj: Any) -> Any:
        """预处理对象（序列化前）"""
        if isinstance(obj, dict):
            return {k: self._preprocess_object(v) for k, v in obj.items()}
        elif isinstance(obj, (list, tuple)):
            return [self._preprocess_object(item) for item in obj]
        elif type(obj) in self.encoder.custom_encoders:
            return self.encoder.encode_object(obj)
        elif hasattr(obj, '__dict__'):
            return self.encoder.encode_object(obj)
        else:
            return obj

    def _update_serialize_time(self, duration: float):
        """更新序列化时间统计"""
        if self.stats.total_serialized == 1:
            self.stats.avg_serialize_time = duration
        else:
            alpha = 0.1
            self.stats.avg_serialize_time = (
                    alpha * duration + (1 - alpha) * self.stats.avg_serialize_time
            )

    def _update_deserialize_time(self, duration: float):
        """更新反序列化时间统计"""
        if self.stats.total_deserialized == 1:
            self.stats.avg_deserialize_time = duration
        else:
            alpha = 0.1
            self.stats.avg_deserialize_time = (
                    alpha * duration + (1 - alpha) * self.stats.avg_deserialize_time
            )

    def get_stats(self) -> Dict[str, Any]:
        """获取序列化统计信息"""
        stats_dict = asdict(self.stats)

        # 计算压缩比
        if self.stats.total_bytes > 0:
            stats_dict['compression_ratio'] = (
                    1.0 - (self.stats.compressed_bytes / self.stats.total_bytes)
            ) if self.stats.compressed_bytes > 0 else 0.0

        # 计算性能指标
        if self.stats.avg_serialize_time > 0:
            stats_dict['serialize_rate'] = 1.0 / self.stats.avg_serialize_time
        if self.stats.avg_deserialize_time > 0:
            stats_dict['deserialize_rate'] = 1.0 / self.stats.avg_deserialize_time

        # 添加配置信息
        stats_dict['config'] = {
            'json_lib': self.json_lib,
            'use_compression': self.use_compression,
            'compression_algorithm': self.compression_algorithm,
            'compression_threshold': self.compression_threshold,
            'ensure_ascii': self.ensure_ascii,
        }

        return stats_dict

    def reset_stats(self):
        """重置统计信息"""
        self.stats = JsonStats()

    def register_encoder(self, type_class, encoder_func):
        """注册自定义编码器"""
        self.encoder.register_encoder(type_class, encoder_func)

    def register_decoder(self, type_name, decoder_func):
        """注册自定义解码器"""
        self.decoder.register_decoder(type_name, decoder_func)

    def benchmark(self, test_data: list, iterations: int = 1000) -> Dict[str, Any]:
        """性能基准测试"""
        import random

        # 重置统计
        self.reset_stats()

        # 准备测试数据
        test_objects = []
        for _ in range(iterations):
            data_item = random.choice(test_data)
            test_objects.append(data_item)

        # 序列化基准测试
        start_time = time.perf_counter()
        serialized_messages = []

        for obj in test_objects:
            serialized = self.dumps(obj)
            serialized_messages.append(serialized)

        serialize_time = time.perf_counter() - start_time

        # 反序列化基准测试
        start_time = time.perf_counter()

        for serialized in serialized_messages:
            self.loads(serialized)

        deserialize_time = time.perf_counter() - start_time

        # 计算结果
        total_size = sum(len(msg) for msg in serialized_messages)

        return {
            'iterations': iterations,
            'serialize_time': serialize_time,
            'deserialize_time': deserialize_time,
            'serialize_rate': iterations / serialize_time,
            'deserialize_rate': iterations / deserialize_time,
            'total_size': total_size,
            'avg_message_size': total_size / iterations,
            'throughput_mbps': (total_size / (1024 * 1024)) / (serialize_time + deserialize_time),
            'stats': self.get_stats()
        }

    def create_message(self, msg_type: str, data: Any, **metadata) -> Dict[str, Any]:
        """创建标准消息格式"""
        return {
            'header': {
                'type': msg_type,
                'timestamp': time.time(),
                'seq_id': int(time.time() * 1000000),
                'correlation_id': str(uuid.uuid4()),
                'metadata': metadata
            },
            'payload': data
        }


# 使用示例
if __name__ == "__main__":
    # 创建序列化器
    serializer = OptimizedJsonSerializer(use_compression=True)

    # 测试基础数据类型
    test_data = {
        'string': 'Hello, JSON!',
        'number': 42,
        'float': 3.14159,
        'boolean': True,
        'null': None,
        'list': [1, 2, 3, 4, 5],
        'dict': {'nested': 'value'},
        'datetime': datetime.now(),
        'date': date.today(),
        'decimal': Decimal('123.456'),
        'bytes': b'binary data',
        'set': {1, 2, 3},
    }

    # 序列化
    serialized = serializer.dumps(test_data)
    print(f"序列化数据: {len(serialized)} bytes")

    # 反序列化
    deserialized = serializer.loads(serialized)
    print(f"反序列化结果: {deserialized}")

    # 测试流式处理
    test_objects = [
        {'id': i, 'data': f'message_{i}', 'timestamp': datetime.now()}
        for i in range(10)
    ]

    stream_serialized = serializer.stream_dumps(test_objects)
    print(f"流式序列化: {len(stream_serialized)} bytes")

    stream_deserialized = serializer.stream_loads(stream_serialized)
    print(f"流式反序列化: {len(stream_deserialized)} objects")

    # 获取统计信息
    stats = serializer.get_stats()
    print(f"统计信息: {stats}")


    # 测试自定义类型
    @dataclass
    class CustomMessage:
        id: int
        content: str
        timestamp: datetime


    # 注册自定义类型
    def encode_custom_message(msg):
        return {
            '__type__': 'custom_message',
            '__value__': {
                'id': msg.id,
                'content': msg.content,
                'timestamp': msg.timestamp.isoformat()
            }
        }


    def decode_custom_message(data):
        return CustomMessage(
            id=data['id'],
            content=data['content'],
            timestamp=datetime.fromisoformat(data['timestamp'])
        )


    serializer.register_encoder(CustomMessage, encode_custom_message)
    serializer.register_decoder('custom_message', decode_custom_message)

    # 测试自定义类型
    custom_msg = CustomMessage(1, "Hello Custom!", datetime.now())
    custom_serialized = serializer.dumps(custom_msg)
    custom_deserialized = serializer.loads(custom_serialized)

    print(f"自定义类型: {custom_deserialized}")

    # 比较不同JSON库的性能
    if len(JSON_LIBS) > 1:
        print("\n性能比较:")
        test_obj = {'data': 'x' * 1000, 'numbers': list(range(100))}

        for lib_name in JSON_LIBS.keys():
            try:
                lib_serializer = OptimizedJsonSerializer(json_lib=lib_name, use_compression=False)

                start_time = time.perf_counter()
                for _ in range(1000):
                    serialized = lib_serializer.dumps(test_obj)
                    lib_serializer.loads(serialized)
                duration = time.perf_counter() - start_time

                print(f"  {lib_name}: {1000 / duration:.0f} ops/s")
            except Exception as e:
                print(f"  {lib_name}: 错误 - {e}")